import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder
import lightgbm as lgb

# 读取数据
df = pd.read_csv("your_data.csv")

# 转换日期格式
df['date'] = pd.to_datetime(df['date'])
df.sort_values(['spu_code', 'date'], inplace=True)

# 编码款号
le = LabelEncoder()
df['spu_code_encoded'] = le.fit_transform(df['spu_code'])


# 2、特征工程
def create_features(df):
    # # 时间特征
    # df['day_of_week'] = df['date'].dt.dayofweek
    # df['month'] = df['date'].dt.month

    # # 历史销售特征
    # df['7d_avg_sale'] = df.groupby('spu_code')['sale_qty'].transform(
    #     lambda x: x.rolling(7, min_periods=1).mean().shift(1))
    #
    # df['14d_avg_sale'] = df.groupby('spu_code')['sale_qty'].transform(
    #     lambda x: x.rolling(14, min_periods=1).mean().shift(1))

    # 上新相关特征
    df['days_since_launch'] = df.groupby('spu_code').cumcount()

    return df


df = create_features(df)


#  3、滚动训练与预测函数
def rolling_forecast(df, train_days=14, forecast_days=14):
    # 按日期排序
    dates = df['date']
    dates.sort()

    # 存储预测结果
    all_predictions = []

    # 滚动窗口：从第31天开始预测
    for i in range(train_days, len(dates)):
        current_date = dates[i]

        # 训练集：过去30天数据
        train_mask = (df['date'] < current_date) & (df['date'] >= dates[i - train_days])
        train_data = df[train_mask]

        # 特征/标签
        features = ['spu_code_encoded', 'days_since_launch', 'before_new_total',
                    'new_sale_qty'
            # , 'day_of_week', 'month', '7d_avg_sale', '14d_avg_sale'
                    ]
        X_train = train_data[features]
        y_train = train_data['sale_qty']

        # 训练模型
        model = lgb.LGBMRegressor(num_leaves=31, learning_rate=0.05, n_estimators=100)
        model.fit(X_train, y_train)

        # 预测未来14天
        future_dates = pd.date_range(current_date, periods=forecast_days + 1)[1:]
        future_predictions = []

        for j, pred_date in enumerate(future_dates, 1):
            # 递归更新特征
            pred_data = df[df['date'] == pred_date].copy()
            if pred_data.empty:
                pred_data = df[df['date'] == current_date].copy()
                pred_data['date'] = pred_date
                pred_data['days_since_launch'] += j

            # 创建预测特征
            pred_data = create_features(pred_data)
            X_pred = pred_data[features]

            # 预测并更新数据
            pred = model.predict(X_pred)
            pred_data['sale_qty'] = pred
            df = pd.concat([df, pred_data], ignore_index=True)
            future_predictions.append(pred_data)

        all_predictions.extend(future_predictions)

    return pd.concat(all_predictions)


# 执行预测
predictions = rolling_forecast(df)


#  4、结果输出
# 提取最后14天预测结果
final_forecast = predictions.groupby('spu_code').tail(14)[
    ['spu_code', 'date', 'sale_qty']]

# 保存结果
final_forecast.to_excel("14_day_forecast.xlsx", index=False)